# 转换注解为yolo
# -*- coding: utf-8 -*-
# @Time    : 2023-12-14 15:47
# @Author  : Jiang Liu
import os
import xml.etree.ElementTree as ET

import cv2
from tqdm import tqdm


def convert_annotation(xml_file, img_file, class_mapping):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    img = cv2.imread(img_file)
    # img.show()
    w = img.shape[1]
    h = img.shape[0]

    yolo_annotations = []
    for obj in root.iter('object'):
        cls = obj.find('name').text
        if cls not in class_mapping:
            print(f"Warning: class '{cls}' not in class_mapping in {xml_file}")
            continue
        cls_id = class_mapping[cls]

        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('xmax').text), float(xmlbox.find('ymax').text))
        bb = convert_to_yolo_format(b, w, h)
        yolo_annotations.append(f"{cls_id} " + " ".join([str(a) for a in bb]))

    return yolo_annotations


def convert_to_yolo_format(box, img_width, img_height):
    dw = 1. / img_width
    dh = 1. / img_height
    x = (box[0] + box[2]) / 2.0 - 1
    y = (box[1] + box[3]) / 2.0 - 1
    w = box[2] - box[0]
    h = box[3] - box[1]
    return [x * dw, y * dh, w * dw, h * dh]


def main():
    # Example class mapping
    class_mapping = {'40070001': 0}  # 根据你的类别调整

    # Directory containing XML files
    xml_dir = './Annotations/'
    img_dir = './JPEGImages'
    txt_dir = './labelsLJ'

    for xml_filename in tqdm(os.listdir(xml_dir)):
        if xml_filename.endswith('.xml'):
            basename = os.path.splitext(xml_filename)[0]
            img_filename = basename + '.jpg'
            yolo_annots = convert_annotation(os.path.join(xml_dir, xml_filename), os.path.join(img_dir, img_filename),
                                             class_mapping)
            txt_filename = basename + '.txt'
            with open(os.path.join(txt_dir, txt_filename), 'w') as f:
                f.write('\n'.join(yolo_annots))


if __name__ == '__main__':
    os.mkdir('labelsLJ')
    main()
